from marshmallow import Schema, fields, validate
from datetime import datetime

class TimeSlotSchema(Schema):
    """时间段模式"""
    startTime = fields.Time(required=True)
    endTime = fields.Time(required=True)
    capacity = fields.Integer(dump_default=10)
    period = fields.Str(validate=validate.OneOf(['上午', '下午', '晚上']), required=True)

class ScheduleSchema(Schema):
    """排班模式"""
    doctorId = fields.Str(required=True)
    departmentId = fields.Str(required=True)
    date = fields.Date(required=True)
    timeSlots = fields.List(fields.Nested(TimeSlotSchema), required=True)
    room = fields.Str()
    registrationType = fields.Str(validate=validate.OneOf(['普通门诊', '专家门诊', '特需门诊']), dump_default='普通门诊')
    registrationFee = fields.Float(dump_default=0.0)
    status = fields.Str(validate=validate.OneOf(['正常', '已满', '停诊']), dump_default='正常')
    remainingCapacity = fields.Dict(keys=fields.Str(), values=fields.Integer())
    remarks = fields.Str()
    createdAt = fields.DateTime(dump_default=datetime.utcnow)
    updatedAt = fields.DateTime(dump_default=datetime.utcnow)

class ScheduleUpdateSchema(Schema):
    """排班更新模式"""
    timeSlots = fields.List(fields.Nested(TimeSlotSchema))
    room = fields.Str()
    registrationType = fields.Str(validate=validate.OneOf(['普通门诊', '专家门诊', '特需门诊']))
    registrationFee = fields.Float()
    status = fields.Str(validate=validate.OneOf(['正常', '已满', '停诊']))
    remainingCapacity = fields.Dict(keys=fields.Str(), values=fields.Integer())
    remarks = fields.Str()

class ScheduleQuerySchema(Schema):
    """排班查询模式"""
    doctorId = fields.Str()
    departmentId = fields.Str()
    startDate = fields.Date()
    endDate = fields.Date()
    registrationType = fields.Str(validate=validate.OneOf(['普通门诊', '专家门诊', '特需门诊', 'all']))
    status = fields.Str(validate=validate.OneOf(['正常', '已满', '停诊', 'all'])) 